/*

Copyright (C) 2000 - 2008 Christian Kreibich <christian@whoop.org>.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to
deal in the Software without restriction, including without limitation the
rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
sell copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies of the Software and its documentation and acknowledgment shall be
given in the documentation and software packages that this Software was
used.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

*/
#ifdef HAVE_CONFIG_H
#  include <config.h>
#endif

#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <glib.h>

#include <nd.h>
#include <nd_protocol.h>
#include <nd_protocol_registry.h>
#include <nd_packet_iterator.h>
#include <nd_gui.h>
#include "nd_tcp.h"
#include "nd_tcb.h"

struct nd_tcb_conn
{
  struct in_addr  ip_src;
  struct in_addr  ip_dst;

  guint16         th_sport;
  guint16         th_dport;

  guint32         src_seq;
  guint32         dst_seq;

  int             first_packet_index;

  gboolean        is_reverse;
};

struct nd_tcb
{
  GHashTable    *table;
};


static ND_TCBConn *
tcb_conn_new(void)
{
  ND_TCBConn *tcbc;

  tcbc = g_new0(ND_TCBConn, 1);
  return tcbc;
}

static void
tcb_conn_free(ND_TCBConn *conn)
{
  if (!conn)
    return;

  g_free(conn);
}


static void
tcb_conn_update_trace(ND_TCBConn *conn, ND_Trace *trace)
{
  ND_PacketIterator     pit;

  gtk_clist_freeze(GTK_CLIST(trace->list));

  for (nd_pit_init_mode(&pit, trace, ND_PACKET_IT_PART); nd_pit_get(&pit); nd_pit_next(&pit))
    {
      if (nd_tcb_is_match(conn, nd_pit_get(&pit)))
	{
	  nd_gui_list_update_packet_state_at_index(nd_pit_get(&pit),
						   nd_pit_get_index(&pit));
	}
    }

  gtk_clist_thaw(GTK_CLIST(trace->list));
}


static guint
tcb_hash_conn(gconstpointer pointer)
{
  const ND_TCBConn *key = (const ND_TCBConn *) pointer;
  
  return
    ((guint32) key->th_sport ^ (guint32) key->ip_src.s_addr)   ^
    ((guint32) key->th_dport ^ (guint32) key->ip_dst.s_addr);
}


static gint
tcb_cmp_conn(gconstpointer p1,
	     gconstpointer p2)
{
  ND_TCBConn *key1 = (ND_TCBConn *) p1;
  ND_TCBConn *key2 = (ND_TCBConn *) p2;

  if (key1->ip_src.s_addr == key2->ip_src.s_addr &&
      key1->ip_dst.s_addr == key2->ip_dst.s_addr &&
      key1->th_sport == key2->th_sport           &&
      key1->th_dport == key2->th_dport)
    {
      key1->is_reverse = FALSE;
      key2->is_reverse = FALSE;
 
      return TRUE;
    }

  if (key1->ip_src.s_addr == key2->ip_dst.s_addr &&
      key1->ip_dst.s_addr == key2->ip_src.s_addr &&
      key1->th_sport == key2->th_dport           &&
      key1->th_dport == key2->th_sport)
    {
      key1->is_reverse = TRUE;
      key2->is_reverse = TRUE;

      return TRUE;
    }

  return FALSE;
}


ND_TCB *
nd_tcb_new(void)
{
  ND_TCB *tcb = g_new0(ND_TCB, 1);
  
  tcb->table = g_hash_table_new(tcb_hash_conn, tcb_cmp_conn);
  return tcb;
}


static void
tcb_free_cb(gpointer key, gpointer val, gpointer user_data)
{
  tcb_conn_free((ND_TCBConn *) val);

  return;
  TOUCH(key);
  TOUCH(user_data);
}


void
nd_tcb_free(ND_TCB *tcb)
{
  if (!tcb)
    return;

  g_hash_table_foreach(tcb->table, tcb_free_cb, NULL);
  g_hash_table_destroy(tcb->table);
  g_free(tcb);
}


ND_TCBConn * 
nd_tcb_lookup(ND_TCB *tcb, const ND_Packet *packet, gboolean *is_reverse)
{
  ND_TCBConn      conn;
  ND_TCBConn     *result;
  struct ip      *iphdr;
  struct tcphdr  *tcphdr;

  if (!tcb || !packet)
    return NULL;

  if (!nd_tcp_get_first(packet, &iphdr, &tcphdr))
    return NULL;

  memset(&conn, 0, sizeof(ND_TCBConn));
  conn.ip_src = iphdr->ip_src;
  conn.ip_dst = iphdr->ip_dst;
  conn.th_sport = tcphdr->th_sport;
  conn.th_dport = tcphdr->th_dport;
  conn.is_reverse = FALSE;

  result = g_hash_table_lookup(tcb->table, &conn);

  if (result)
    {
      if (is_reverse)
	*is_reverse = conn.is_reverse;
      
      return result;
    }
  
  return NULL;
}


void        
nd_tcb_update(ND_TCB *tcb, const ND_Packet *packet, int index)
{
  struct tcphdr  *tcphdr;
  struct ip      *iphdr;
  ND_TCBConn     *tcbc;
  gboolean        is_reverse = FALSE;
  gboolean        update_trace = FALSE;

  if (!tcb || !packet)
    return;

  if (!nd_tcp_get_first(packet, &iphdr, &tcphdr))
    return;

  if (! (tcbc = nd_tcb_lookup(tcb, packet, &is_reverse)))
    {
      tcbc = tcb_conn_new();

      tcbc->ip_src  = iphdr->ip_src;
      tcbc->ip_dst  = iphdr->ip_dst;
      tcbc->src_seq = ntohl(tcphdr->th_seq);
      
      if (ntohl(tcphdr->th_ack) != 0)
	tcbc->dst_seq = ntohl(tcphdr->th_ack) - 1; 
      
      tcbc->th_sport = tcphdr->th_sport;
      tcbc->th_dport = tcphdr->th_dport;
      tcbc->first_packet_index = nd_packet_get_index(packet);
      
      D(("TCB initialized to index %i for %u %u\n",
	 tcbc->first_packet_index, tcbc->src_seq, tcbc->dst_seq));
      
      g_hash_table_insert(tcb->table, tcbc, tcbc);
    }
  else
    {
      /* I understand tcpdump's TCP state management as follows:
       *
       * - the first packet of a tcpdump connection is used as
       *   the basis for the relative seq/ack calculations.
       * - these numbers are never changed. Should a packet be
       *   processed that has a smaller seq/ack than the one
       *   in the state structure, the relative number calculation
       *   would become negative but due to u_... types is
       *   interpreted as a positive number.
       *
       * We will emulate this behaviour here and never print
       * negative numbers. We will also only update the TCB when
       * a packet of the same connection is moved before the
       * previously first packet of a connection.
       */

      if (!is_reverse)
	{
	  if (ntohl(tcphdr->th_seq) != tcbc->src_seq)
	    {
	      if (index < 0)
		index = nd_packet_get_index(packet);

	      if (index <= tcbc->first_packet_index)
		{
		  tcbc->first_packet_index = index;
		  tcbc->src_seq = ntohl(tcphdr->th_seq);

		  update_trace = TRUE;

		  D(("TCB updated to index %i for %u %u\n",
		     tcbc->first_packet_index, tcbc->src_seq, tcbc->dst_seq));
		}
	    }
	  
	  if (tcbc->dst_seq == 0 && ntohl(tcphdr->th_ack) != 0)
	    {
	      tcbc->dst_seq = ntohl(tcphdr->th_ack) - 1;
	      return;
	    }

	  if (ntohl(tcphdr->th_ack) != 0 &&
	      ntohl(tcphdr->th_ack) - 1 != tcbc->dst_seq)
	    {
	      if (index < 0)
		index = nd_packet_get_index(packet);

	      if (index <= tcbc->first_packet_index)
		{
		  tcbc->first_packet_index = index;
		  tcbc->dst_seq = ntohl(tcphdr->th_ack) - 1; 
		  
		  update_trace = TRUE;

		  D(("TCB updated to index %i for %u %u\n",
		     tcbc->first_packet_index, tcbc->src_seq, tcbc->dst_seq));
		}
	    }
	}
      else
	{
	  if (ntohl(tcphdr->th_seq) != tcbc->dst_seq)
	    {
	      if (index < 0)
		index = nd_packet_get_index(packet);

	      if (index <= tcbc->first_packet_index)
		{
		  tcbc->first_packet_index = index;
		  tcbc->dst_seq = ntohl(tcphdr->th_seq);
		  
		  update_trace = TRUE;

		  D(("TCB updated rev to index %i for %u %u\n",
		     tcbc->first_packet_index, tcbc->src_seq, tcbc->dst_seq));
		}
	    }
	  
	  if (tcbc->src_seq == 0 && ntohl(tcphdr->th_ack) != 0)
	    {
	      tcbc->src_seq = ntohl(tcphdr->th_ack) - 1;
	      return;
	    }

	  if (ntohl(tcphdr->th_ack) != 0 &&
	      ntohl(tcphdr->th_ack) - 1 != tcbc->dst_seq)
	    {
	      if (index < 0)
		index = nd_packet_get_index(packet);

	      if (index <= tcbc->first_packet_index)
		{
		  tcbc->first_packet_index = index;
		  tcbc->src_seq = ntohl(tcphdr->th_ack) - 1; 
		  
		  update_trace = TRUE;

		  D(("TCB updated rev to index %i for %u %u\n",
		     tcbc->first_packet_index, tcbc->src_seq, tcbc->dst_seq));
		}
	    }
	}

      if (nd_tcp_get_state_mode() == ND_TCP_STATE_UPDATE_ALL && update_trace)
	tcb_conn_update_trace(tcbc, packet->trace);
    }
}


gboolean        
nd_tcb_is_match(const ND_TCBConn *conn,
		const ND_Packet *packet)
{
  struct tcphdr  *tcphdr;
  struct ip      *iphdr;

  if (!conn || !packet)
    return FALSE;

  if (!nd_tcp_get_first(packet, &iphdr, &tcphdr))
    return FALSE;
  
  if (iphdr->ip_src.s_addr == conn->ip_src.s_addr &&
      iphdr->ip_dst.s_addr == conn->ip_dst.s_addr &&
      tcphdr->th_sport == conn->th_sport          &&
      tcphdr->th_dport == conn->th_dport)
    return TRUE;

  if (iphdr->ip_src.s_addr == conn->ip_dst.s_addr &&
      iphdr->ip_dst.s_addr == conn->ip_src.s_addr &&
      tcphdr->th_sport == conn->th_dport          &&
      tcphdr->th_dport == conn->th_sport)
    return TRUE;

  return FALSE;
}


gboolean        
nd_tcb_conn_reverse(const ND_TCBConn *conn)
{
  return conn->is_reverse;
}


gboolean
nd_tcb_conn_recv_known(ND_TCBConn *tcbc)
{
  if (!tcbc || tcbc->dst_seq == 0)
    return FALSE;
  
  return TRUE;
}


gboolean
nd_tcb_conn_get_rel_seq(const ND_TCBConn *tcbc,
			const struct ip *iphdr, const struct tcphdr *tcphdr,
			guint32 *seq_start, guint32 *seq_end)
{
  int      length;
  guint32  seq;
  
  if (!tcbc || !tcphdr || !iphdr || !seq_start || !seq_end)
    return FALSE;
  
  length = ntohs(iphdr->ip_len) - (iphdr->ip_hl * 4) -
    (tcphdr->th_off * 4);
  
  if (iphdr->ip_src.s_addr == tcbc->ip_src.s_addr)
    {
      seq = ntohl(tcphdr->th_seq);

      if (tcbc->src_seq == seq)
	{
	  *seq_start = tcbc->src_seq;
	  *seq_end   = tcbc->src_seq + length;
	  return FALSE;
	}
      else
	{
	  *seq_start = seq - tcbc->src_seq;
	  *seq_end   = seq - tcbc->src_seq + length;
	  return TRUE;
	}
    }
  else if (iphdr->ip_src.s_addr == tcbc->ip_dst.s_addr)
    {
      seq = ntohl(tcphdr->th_seq);

      if (tcbc->dst_seq == seq)
	{
	  *seq_start = tcbc->dst_seq;
	  *seq_end   = tcbc->dst_seq + length;
	  return FALSE;
	}
      else
	{
	  *seq_start = seq - tcbc->dst_seq;
	  *seq_end   = seq - tcbc->dst_seq + length;
	  return TRUE;
	}
    }

  D(("Mismatch!\n"));
  return FALSE;
}


gboolean
nd_tcb_conn_get_rel_ack(const ND_TCBConn *tcbc,
			const struct ip *iphdr, const struct tcphdr *tcphdr,
			gboolean force_rel, guint32 *ack)
{
  guint32 delta;

  if (!ack)
    return FALSE;
  
  if (!tcbc || !iphdr || !tcphdr)
    {
      *ack = 0;
      return FALSE;
    }

  if (iphdr->ip_src.s_addr == tcbc->ip_src.s_addr)
    {
      delta = (ntohl(tcphdr->th_ack) - tcbc->dst_seq);
      if (delta == 1 && !force_rel)
	{
	  *ack = (tcbc->dst_seq + 1); 
	  return FALSE;
	}

      *ack = delta;
      return TRUE;
    }
  else if (iphdr->ip_src.s_addr == tcbc->ip_dst.s_addr)
    {
      delta = (ntohl(tcphdr->th_ack) - tcbc->src_seq);
      if (delta == 1 && !force_rel)
	{
	  *ack = (tcbc->src_seq + 1);
	  return FALSE;
	}

      *ack = delta;
      return TRUE;
    }
  
  D(("Mismatch!\n"));
  return FALSE;
}
